import torch
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch import nn

x = torch.linspace(-5, 5, 20)
x = torch.sort(x).values
x = x.reshape(-1, 1)
y = x * torch.cos(x)

model = nn.Sequential(
    nn.Linear(1, 20),
    nn.Tanh(),
    nn.Linear(20, 1)
)

criterion = nn.MSELoss()
sgd = torch.optim.Adam(model.parameters(), lr=0.03)

model.train()
epochs = 100
for epoch in range(epochs):
    sgd.zero_grad()
    y_predict = model(x)
    loss = criterion(y_predict, y)
    loss.backward()
    sgd.step()
    if epoch % 100 == 0:
        print(f"epoch {epoch + 1} / {epochs} -- loss: {loss.item():.2f}")

model.eval()
plt.plot(x, y, 'r-')
y_predict = model(x)
plt.plot(x.detach().numpy(), y_predict.detach().numpy(), 'b--')
plt.show()
